import json
import logging
from contextlib import AsyncExitStack

from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
from langgraph.graph import END, START, StateGraph
from langgraph.prebuilt import ToolNode, tools_condition

from prompt import agent_prompt_template
from state import MessagesState
from tool import (
    get_attractions_information,
    get_coordinate,
    get_static_map,
    nearby,
    route_planning,
    web_search,
)

logger = logging.getLogger(__name__)


def create_graph(tools, model_provider: str = None, is_async=True):
    graph = StateGraph(MessagesState)

    if is_async:
        from agent import AsyncAgent as Agent
    else:
        from agent import SyncAgent as Agent

    agent = Agent(
        tools=tools, prompt_template=agent_prompt_template, model_provider=model_provider, temperature=0.2
    )

    graph.add_node("agent", agent)
    graph.add_node("tools", ToolNode(tools))

    graph.add_edge(START, "agent")
    graph.add_conditional_edges("agent", tools_condition)
    graph.add_edge("tools", "agent")
    graph.add_edge("agent", END)

    return graph


async def init_app(stack: AsyncExitStack, model_name: str = None, is_async=True):
    tools = [get_static_map, get_coordinate, get_attractions_information, web_search, route_planning, nearby]
    graph = create_graph(tools, model_name, is_async)

    sqlite_saver = await stack.enter_async_context(
        AsyncSqliteSaver.from_conn_string(r"D:\code\flying_pig\static\checkpoint\checkpoint.sqlite")
    )

    app = graph.compile(checkpointer=sqlite_saver)

    return app


def get_mcp_servers(file_path: str = r"D:\code\flying_pig\mcp_servers_config.json") -> dict:
    with open(file_path, "r", encoding="utf-8") as f:
        return json.load(f).get("mcpServers", {})


async def init_app_withmcp(stack: AsyncExitStack, model_name: str = None, is_async=True):
    from langchain_mcp_adapters.client import MultiServerMCPClient

    mcp_client = MultiServerMCPClient(get_mcp_servers())

    tools = await mcp_client.get_tools()
    logger.info(f"一共获取到 {len(tools)} 个工具")

    graph = create_graph(tools, model_name, is_async)

    sqlite_saver = await stack.enter_async_context(
        AsyncSqliteSaver.from_conn_string(r"D:\code\flying_pig\static\checkpoint\checkpoint.sqlite")
    )

    app = graph.compile(checkpointer=sqlite_saver)

    return app
